from PyQt5.QtWidgets import QApplication,QMainWindow,QGridLayout
from PyQt5.QtCore import QTimer
from PyQt5.QtCore import QSettings
from matplotlib.pyplot import ylabel
from numpy import arange

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from pandas import read_csv
import pypianoroll

class Figure_Canvas_EEG(FigureCanvas):
    def __init__(self,parent=None,width=3.9,height=2.7,dpi=100):
        self.fig = Figure(figsize=(width,height),dpi=100)
        # self.fig.patch.set_facecolor('#323232')
        super(Figure_Canvas_EEG,self).__init__(self.fig)
        self.ax = self.fig.add_subplot(111)
        # self.ax.patch.set_facecolor('#323232')

    def reset_axis(self):
        self.ax.cla()
        # self.ax.patch.set_facecolor('#323232')
        self.ax.set_ylim(-1500, 1500)
        self.ax.set_xlabel("Time Step (1 step = 0.25 s)")

class Plot_EEG():
    def __init__(self, parent, sample_step=250):
        self.widget = parent
        self.sample_step = sample_step
        self.txt_path = ""
        self.LineFigure = Figure_Canvas_EEG()
        self.LineFigureLayout = QGridLayout(self.widget) 
        self.LineFigureLayout.addWidget(self.LineFigure)
        self.mapping = {
            'Ch1': 'Fp1',
            'Ch2': 'Fp2',
            'Ch3': 'Fz',
            'Ch4': 'C3',
            'Ch5': 'C4',
            'Ch6': 'Pz',
            'Ch7': 'O1',
            'Ch8': 'O2'
        }
        
    def set_path(self, txt_path):
        self.txt_path = txt_path
        self.data = read_csv(self.txt_path, sep='\t', header=1)

    def change_channel(self, channel):
        self.channel = channel
        ax = self.LineFigure.ax
        if self.txt_path != "":
            self.LineFigure.reset_axis()
            # ax.axes.yaxis.set_visible(False)
            ax.set_ylabel("Channels")
            level = 0
            ylabel = []
            offset = []
            for c in channel:
                col = "Ch" + str(c) 
                y = self.data[col][0::self.sample_step] + level
                offset.append(y[0])
                level += 1000
                x = arange(0, len(y))
                ax.plot(x,y,linewidth=0.5)
                ax.set_xlim(0, len(y))
                ylabel.append(self.mapping[col])

            # offset = list(range(0, 1000 * (len(channel)), 1000))  等距的实现方式
            ax.set_yticks(offset, ylabel)
            ax.autoscale(enable=True, axis='y', tight=None)
            self.LineFigure.fig.canvas.draw() 
            self.LineFigure.fig.canvas.flush_events() 


class Plot_EEG32():
    def __init__(self, parent, sample_step=250):
        self.widget = parent
        self.sample_step = sample_step
        self.txt_path = ""
        self.LineFigure = Figure_Canvas_EEG()
        self.LineFigureLayout = QGridLayout(self.widget) 
        self.LineFigureLayout.addWidget(self.LineFigure)
        
    def set_path(self, txt_path):
        self.txt_path = txt_path
        self.data = read_csv(self.txt_path, sep='\t', header=None, skiprows=1)
        # print(self.data.head())
        # print(self.data.shape)

    def change_channel(self, channel):
        self.channel = channel
        ax = self.LineFigure.ax
        if self.txt_path != "":
            self.LineFigure.reset_axis()
            ax.axes.yaxis.set_visible(False)
            ax.set_ylabel("Channels")
            self.LineFigure.ax.set_xlabel("Time")
            level = 0
            for c in channel:
                y = self.data.iloc[0::self.sample_step, c] + level
                level += 1500
                x = arange(0, len(y))
                ax.plot(x,y,linewidth=0.5)
                ax.set_xlim(0, len(y))

            ax.autoscale(enable=True, axis='y', tight=None)
            self.LineFigure.fig.canvas.draw()
            self.LineFigure.fig.canvas.flush_events()

class Figure_Canvas_Music(FigureCanvas):
    def __init__(self,parent=None,width=3.9,height=2.7,dpi=100):
        self.fig = Figure(figsize=(width,height),dpi=100)
        # self.fig.patch.set_facecolor('#323232')
        super(Figure_Canvas_Music,self).__init__(self.fig)
        self.ax = self.fig.add_subplot(111)
        # self.ax.patch.set_facecolor('#323232')

    def reset_axis(self):
        self.ax.cla()
        # self.ax.patch.set_facecolor('#323232')
        self.ax.set_ylim(-1500, 1500)

class Plot_Music():
    def __init__(self, parent, channel=1 ,sample_step=250):
        self.widget = parent
        self.channel = channel
        self.sample_step = sample_step
        self.settings = QSettings('Config.ini', QSettings.IniFormat)
        self.txt_path = self.settings.value('General/musicpath')
        self.LineFigure = Figure_Canvas_Music()
        self.LineFigureLayout = QGridLayout(self.widget) 
        self.LineFigureLayout.addWidget(self.LineFigure)
        
    def plot(self):
        multitrack = pypianoroll.read(self.txt_path)
        multitrack.trim(0, 12 * multitrack.resolution)
        multitrack.binarize()
        pypianoroll.plot_multitrack(multitrack, [self.LineFigure.ax])
        self.LineFigure.fig.canvas.draw() 
        self.LineFigure.fig.canvas.flush_events() 
